import os.path

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_core.prompt_values import PromptValue
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, BasePromptTemplate

from infrastructure.constant import file


class PromptTemplateService:
    @staticmethod
    def from_message_list_with_system_message(messages: list[dict], system_prompt_content=None) -> PromptValue:
        if not system_prompt_content:
            system_prompt_content = PromptTemplateService.get_string_prompt_with_system_message()
        return ChatPromptTemplate.from_messages([
            SystemMessage(content=system_prompt_content),
            *[
                HumanMessage(message.get('content'))
                if message.get('role') in ('user', 'human')
                else AIMessage(message.get('content'))
                for message in messages
            ]
        ]).invoke({})

    @staticmethod
    def get_string_prompt_with_system_message() -> str:
        with open(os.path.join(file.PROMPT_TEMPLATE_PATH, "system_prompt.md"), 'r', encoding='utf-8') as f:
            return f.read().strip()

    @staticmethod
    def get_string_prompt_with_system_message_excel(version: str = None) -> str:
        if not version:
            version = 'v1'
        file_path = f"system_prompt_excel_{version}.md"
        with open(os.path.join(file.PROMPT_TEMPLATE_PATH, file_path), 'r', encoding='utf-8') as f:
            return f.read().strip()
